import copy
import os
import random

import numpy as np
import torch
from nuimages import NuImages
from PIL import Image
from torch.utils.data import Dataset
from torchvision.transforms import InterpolationMode, RandomResizedCrop
from torchvision.transforms.functional import hflip, resize, resized_crop

# import cv2


def collate_pair_fn(list_data):
    """
    Collate function adapted for creating batches with MinkowskiEngine.
    """
    (im, semantic_mask) = list(zip(*list_data))

    im_batch = torch.stack(im)
    msk_batch = torch.stack(semantic_mask).long()

    return {
        "ims": im_batch,
        "msks": msk_batch,
    }


class NuImagesDataset(Dataset):
    """
    Dataset matching a 3D points cloud and an image using projection.
    """

    def __init__(
        self,
        phase,
        config,
        **kwargs,
    ):
        self.phase = phase
        self.nuim = NuImages(
            dataroot="datasets/nuimages", version="v1.0-" + self.phase, verbose=False
        )

        self.config = config

        self.samples = self.nuim.sample

        # idx = np.random.choice(len(self.samples), 5000, replace=False)
        # self.samples = np.array(self.samples)[idx]

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        key_camera_token = sample["key_camera_token"]
        semantic_mask, instance_mask = self.nuim.get_segmentation(key_camera_token)
        semantic_mask = Image.fromarray(semantic_mask)
        sample_data = self.nuim.get("sample_data", key_camera_token)
        im_path = os.path.join(self.nuim.dataroot, sample_data["filename"])
        im = Image.open(im_path)

        i, j, h, w = RandomResizedCrop.get_params(
            im, [0.3, 1.0], self.config["crop_ratio"]
        )
        im = resized_crop(
            im, i, j, h, w, self.config["crop_size"], InterpolationMode.BILINEAR
        )
        semantic_mask = resized_crop(
            semantic_mask,
            i,
            j,
            h,
            w,
            self.config["crop_size"],
            InterpolationMode.NEAREST,
        )

        if self.phase == "train":
            if random.random() < 0.5:
                im = hflip(im)
                semantic_mask = hflip(semantic_mask)

        im = torch.tensor(np.array(im).transpose(2, 0, 1)) / 255
        semantic_mask = torch.tensor(np.array(semantic_mask), dtype=torch.int32)

        """
        nuImages labels:                             eval labels:
        0   "noise"                                  0  "noise"
        1	"animal"                                 1	"barrier",
        2	"human.pedestrian.adult"                 2	"bicycle",
        3	"human.pedestrian.child"                 3	"bus",
        4	"human.pedestrian.construction_worker"   4	"car",
        5	"human.pedestrian.personal_mobility"     5	"construction_vehicle",
        6	"human.pedestrian.police_officer"        6	"motorcycle",
        7	"human.pedestrian.stroller"              7	"pedestrian",
        8	"human.pedestrian.wheelchair"            8	"traffic_cone",
        9	"movable_object.barrier"                 9	"trailer",
        10	"movable_object.debris"                  10	"truck",
        11	"movable_object.pushable_pullable"       11	"driveable_surface",
        12	"movable_object.trafficcone"             12	"other_flat",
        13	"static_object.bicycle_rack"             13	"sidewalk",
        14	"vehicle.bicycle"                        14	"terrain",
        15	"vehicle.bus.bendy"                      15	"manmade",
        16	"vehicle.bus.rigid"                      16	"vegetation",
        17	"vehicle.car"
        18	"vehicle.construction"
        19	"vehicle.emergency.ambulance"
        20	"vehicle.emergency.police"
        21	"vehicle.motorcycle"
        22	"vehicle.trailer"
        23	"vehicle.truck"
        24	'flat.driveable_surface'
        31	'vehicle.ego'
        """

        # labels' names lookup table
        self.eval_labels = {
            0: 0,
            1: 0,
            2: 7,
            3: 7,
            4: 7,
            5: 0,
            6: 7,
            7: 0,
            8: 0,
            9: 1,
            10: 0,
            11: 0,
            12: 8,
            13: 0,
            14: 2,
            15: 3,
            16: 3,
            17: 4,
            18: 5,
            19: 0,
            20: 0,
            21: 6,
            22: 9,
            23: 10,
            24: 11,
            25: 12,
            26: 13,
            27: 14,
            28: 15,
            29: 0,
            30: 16,
            31: 0,
        }  # 25: 12, 26: 13, 27: 14, 28: 15, 29: 0, 30: 16 are not appeared
        semantic_mask = torch.tensor(
            np.vectorize(self.eval_labels.__getitem__)(semantic_mask),
            dtype=torch.int32,
        )

        # tmp_msk = np.array(semantic_mask, dtype=np.uint8)*15
        # tmp_msk = cv2.applyColorMap(tmp_msk, cv2.COLORMAP_RAINBOW)
        # cv2.imwrite("msk.png", tmp_msk)
        # tmp_im = np.array(np.array(im).transpose(1, 2, 0)*255, dtype=np.uint8)
        # cv2.imwrite("tmp_im.png", tmp_im)

        return (
            im,
            semantic_mask,
        )
